package com.example.chatrecommendation.service.impl;

import com.example.chatrecommendation.service.AiPaintingService;
import com.example.chatrecommendation.tool.Translation;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import okhttp3.HttpUrl;
import org.apache.http.client.methods.CloseableHttpResponse;
import org.apache.http.client.methods.HttpUriRequest;
import org.apache.http.client.methods.RequestBuilder;
import org.apache.http.entity.ContentType;
import org.apache.http.entity.StringEntity;
import org.apache.http.impl.client.CloseableHttpClient;
import org.apache.http.impl.client.HttpClients;
import org.apache.http.util.EntityUtils;
import org.json.JSONArray;
import org.json.JSONObject;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.HttpEntity;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.stereotype.Service;
import org.springframework.web.client.RestTemplate;

import javax.crypto.Mac;
import javax.crypto.spec.SecretKeySpec;
import java.net.URI;
import java.net.URL;
import java.nio.charset.StandardCharsets;
import java.text.SimpleDateFormat;
import java.util.*;
import java.util.concurrent.ThreadLocalRandom;


@Service
public class AiPaintingServiceImpl implements AiPaintingService {

    public static final String hostUrl = "https://spark-api.cn-huabei-1.xf-yun.com/v2.1/tti";
    public static final String appid = "3de97f28";
    public static final String apiSecret = "ZGY2NWM1YWMxNGYwYzMxNzg4Y2FjMzZk";
    public static final String apiKey = "cde6d8dd11e0b3b229d47cde71745ad9";
    private static final Logger log = LoggerFactory.getLogger(AiPaintingServiceImpl.class);
    private static RestTemplate restTemplate = new RestTemplate();

    @Autowired
    public AiPaintingServiceImpl(RestTemplate restTemplate) {
        this.restTemplate = restTemplate;
    }


    /** 从响应体中提取Base64字符串
     *
     * @param responseBody
     * @return
     */
    private static String extractBase64FromResponse(String responseBody,String key) {
        ObjectMapper objectMapper = new ObjectMapper();
        try {
            JsonNode jsonNode = objectMapper.readTree(responseBody);
            // 对于spark接口，我们需要从payload.choices.text数组中提取
            if ("spark".equals(key)) {
                JsonNode textArray = jsonNode.path("payload").path("choices").path("text");
                if (textArray.isArray() && textArray.size() > 0) {
                    JsonNode contentNode = textArray.get(0).path("content");
                    if (contentNode.isTextual()) {
                        return contentNode.asText();
                    }
                }
            }
            // 对于sdApi接口，我们直接从images数组中提取
            else if ("sdapi".equals(key)) {
                JsonNode rootNode = objectMapper.readTree(responseBody);
                // 从 JsonNode 对象中提取 "images" 数组
                JsonNode imagesNode = rootNode.get("images");

                // 验证 "images" 是否存在并且是一个数组
                if (imagesNode != null && imagesNode.isArray() && !imagesNode.isEmpty()) {
                    // 提取 "images" 数组中的第一个元素
                    JsonNode firstImageNode = imagesNode.get(0);

                    // 如果存在且为字符串类型，则返回该字符串
                    if (firstImageNode != null && firstImageNode.isTextual()) {
                        return firstImageNode.asText();
                    }
                }
            }
        } catch (Exception e) {
            // 处理异常，例如JSON解析错误
            e.printStackTrace();
        }
        return null; // 如果没有找到或解析失败，返回null
    }

    /**
     * 生成徽章
     * @return
     */
    public String generateHonorBadge() {
        log.info("SD开始生成徽章");

        // 构建请求体
        // seed 1951701925 2287461704 效果最好
        String[] seeds = {"1951701925", "2287461705", "4127465782","3647933624", "3440542522", "3440542523", "3440542525", "1832680489", "3492541044"};
        String randomSeed = seeds[ThreadLocalRandom.current().nextInt(0, seeds.length)];
        log.info("使用的徽章种子：{}", randomSeed);
        String requestJson = """
                {
                  "prompt": "<lora:badge:1>(masterpiece),(best quality),(ultra-detailed),badge,An imaginative and prestigious achievement badge for completing daily tasks consecutively for30 days,featuring a phoenix risirsing from ashes,a calendar,anda trophy.The badge has an ethereal glowing effect with a silver sheen, high-resolution details,and a color scheme of red and gold,symbolizing perseverance and triumph.It's designed with a mythical and victorious style, representing the transfor mative journey of consistent dedication an dsuccess. ",
                  "negative_prompt": "(painting by bad-artist-anime:0.9), (painting by bad-artist:0.9), watermark, text, error, blurry, jpeg artifacts, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, artist name, (worst quality, low quality:1.4), bad anatomy",
                  "cfg_scale": 7,
                  "height": 256,
                  "width": 256,
                  "steps": 30,
                  "sampler_index": "DPM++ 2M Karras",
                  "seed": """ + randomSeed +
                """
                        }
                        """;

        return generateImgs(requestJson);
    }

    /**
     * 生成图像
     * @param prompt
     * @return
     */
    public String generateImage(String prompt) {
        log.info("SD开始生成图像");
        //预处理提示词，删除prompt中开头的
        String preprocessedPrompt = prompt.replaceAll("\\s+", " ");

        if (preprocessedPrompt.contains("<lora:")){
            return generateHonorBadge();
        }

        //翻译提示词
        try {
            Translation translation = new Translation();
            String jsonString = translation.translate(preprocessedPrompt);
            JSONObject jsonObject = new JSONObject(jsonString);
            JSONArray transResults = jsonObject.getJSONArray("trans_result");
            if (transResults.length() > 0) {
                prompt = transResults.getJSONObject(0).getString("dst");
            }
            log.debug("翻译后的prompt：{}", prompt);
        } catch (Exception e) {
            log.error("翻译失败", e);
        }

        // 构建请求体
        String requestJson = """
                {
                  "prompt":"(masterpiece),(best quality),(ultra-detailed),""" + prompt + """
                ",
                  "negative_prompt": "(painting by bad-artist-anime:0.9), (painting by bad-artist:0.9), watermark, text, error, blurry, jpeg artifacts, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, artist name, (worst quality, low quality:1.4), bad anatomy",
                  "cfg_scale": 7.5,
                  "steps": 20,
                  "sampler_index": "DPM++ 2M Karras",
                  "restore_faces": true,
                  "seed": -1
                }
                """;
        log.debug("请求体：\n{}", requestJson);

        return generateImgs(requestJson);
    }

    /**
     * 向ai绘画后端发送请求
     * @param requestJson
     * @return
     */
    private String generateImgs(String requestJson) {
        // 设置请求头
        HttpHeaders headers = new HttpHeaders();
        headers.setContentType(MediaType.APPLICATION_JSON);

        // 创建HttpEntity
        HttpEntity<String> entity = new HttpEntity<>(requestJson, headers);

        ResponseEntity<String> response = restTemplate.postForEntity("http://127.0.0.1:7860/sdapi/v1/txt2img", entity, String.class);

        // 这里应该是从response.getBody()中提取Base64字符串
        String imageBase64 = extractBase64FromResponse(response.getBody(), "sdapi");

        return imageBase64;
    }

    public String sparkAiPainting(String prompt) throws Exception {
        log.info("spark开始生成图像");
        String authUrl = getAuthUrl(hostUrl, apiKey, apiSecret);
        // URL地址正确
        log.debug("authUrl: {}", authUrl);
        String json = "{\n" +
                "  \"header\": {\n" +
                "    \"app_id\": \"" + appid + "\",\n" +
                "    \"uid\": \"" + UUID.randomUUID().toString().substring(0, 15) + "\"\n" +
                "  },\n" +
                "  \"parameter\": {\n" +
                "    \"chat\": {\n" +
                "      \"domain\": \"s291394db\",\n" +
                "      \"temperature\": 0.5,\n" +
                "      \"max_tokens\": 4096,\n" +
                "      \"width\": 512,\n" +
                "      \"height\": 512\n" +
                "    }\n" +
                "  },\n" +
                "  \"payload\": {\n" +
                "    \"message\": {\n" +
                "      \"text\": [\n" +
                "        {\n" +
                "          \"role\": \"user\",\n" +
                "          \"content\": \"" + prompt + "\"\n" +
                "        }\n" +
                "      ]\n" +
                "    }\n" +
                "  }\n" +
                "}";
        // 发起Post请求
        log.debug("ai绘画请求体：{}", json);
        String res = doPostJson(authUrl, null, json);
        String imageBase64 = extractBase64FromResponse(res, "spark");
        return imageBase64;
    }

    public String doPostJson(String url, Map<String, String> urlParams, String json) {
        CloseableHttpClient closeableHttpClient = HttpClients.createDefault();
        CloseableHttpResponse closeableHttpResponse = null;
        String resultString = "";
        try {
            // 创建Http Post请求
            String asciiUrl = URI.create(url).toASCIIString();
            RequestBuilder builder = RequestBuilder.post(asciiUrl);
            builder.setCharset(StandardCharsets.UTF_8);
            if (urlParams != null) {
                for (Map.Entry<String, String> entry : urlParams.entrySet()) {
                    builder.addParameter(entry.getKey(), entry.getValue());
                }
            }
            // 创建请求内容
            StringEntity entity = new StringEntity(json, ContentType.APPLICATION_JSON);
            builder.setEntity(entity);
            HttpUriRequest request = builder.build();
            // 执行http请求
            closeableHttpResponse = closeableHttpClient.execute(request);
            resultString = EntityUtils.toString(closeableHttpResponse.getEntity(), StandardCharsets.UTF_8);
        } catch (Exception e) {
            e.printStackTrace();
        } finally {
            try {
                if (closeableHttpResponse != null) {
                    closeableHttpResponse.close();
                }
                if (closeableHttpClient != null) {
                    closeableHttpClient.close();
                }
            } catch (Exception e) {
                e.printStackTrace();
            }
        }
        return resultString;
    }

    // 鉴权方法
    public String getAuthUrl(String hostUrl, String apiKey, String apiSecret) throws Exception {
        URL url = new URL(hostUrl);
        // 时间
        SimpleDateFormat format = new SimpleDateFormat("EEE, dd MMM yyyy HH:mm:ss z", Locale.US);
        format.setTimeZone(TimeZone.getTimeZone("GMT"));
        String date = format.format(new Date());
        // date="Thu, 12 Oct 2023 03:05:28 GMT";
        // 拼接
        String preStr = "host: " + url.getHost() + "\n" + "date: " + date + "\n" + "POST " + url.getPath() + " HTTP/1.1";
        // System.err.println(preStr);
        // SHA256加密
        Mac mac = Mac.getInstance("hmacsha256");
        SecretKeySpec spec = new SecretKeySpec(apiSecret.getBytes(StandardCharsets.UTF_8), "hmacsha256");
        mac.init(spec);

        byte[] hexDigits = mac.doFinal(preStr.getBytes(StandardCharsets.UTF_8));
        // Base64加密
        String sha = Base64.getEncoder().encodeToString(hexDigits);
        // System.err.println(sha);
        // 拼接
        String authorization = String.format("api_key=\"%s\", algorithm=\"%s\", headers=\"%s\", signature=\"%s\"", apiKey, "hmac-sha256", "host date request-line", sha);
        // 拼接地址
        HttpUrl httpUrl = Objects.requireNonNull(HttpUrl.parse("https://" + url.getHost() + url.getPath())).newBuilder().//
                addQueryParameter("authorization", Base64.getEncoder().encodeToString(authorization.getBytes(StandardCharsets.UTF_8))).//
                addQueryParameter("date", date).//
                addQueryParameter("host", url.getHost()).//
                build();

        // System.err.println(httpUrl.toString());
        return httpUrl.toString();
    }

    @Override
    public String getImage(String prompt) {
        return "";
    }
}
